import os.path
import struct
import random
import hashlib
import uuid
import datetime
from unittest import TestCase

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes, hmac, serialization
from cryptography.hazmat.primitives.asymmetric import ec, rsa, padding
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.x509.oid import NameOID
import socket

from ssl.common import generate_gmt_time, derive_keys


def write_string_to_file(content, file_path, mode='w', is_multi_lines=False):
    try:
        with open(file_path, mode, encoding='utf-8') as file:
            if is_multi_lines:
                for line in content:
                    file.write(line + '\n')
            else:
                file.write(content)
        if mode == 'w':
            print(f"写入模式：成功将内容写入 {file_path}")
        elif mode == 'a':
            print(f"追加模式：成功将内容追加到 {file_path}")
    except Exception as e:
        print(f"写入文件时出现错误: {e}")


g_certificate = ""


# 生成 ServerHello 报文
def create_server_hello(client_hello):
    client_version = client_hello[5:7]
    cipher_suites = client_hello[34:34 + 2 + 2 * int.from_bytes(client_hello[32:34], 'big')]

    handshake_type = b'\x02'  # ServerHello
    client_version = b'\x03\x03'
    server_time = generate_gmt_time()
    server_hello = (
            client_version +  # Negotiated version
            server_time +
            b'\x20' +  # Session ID length (0x00 for new session)
            server_time +  # Session ID
            struct.pack('>H', 0xC030) +  # Selected cipher suite (ECDHE-RSA-AES256-GCM-SHA384)
            b'\x00' +  # Compression method
            b'\x00\x00'  # Extensions (empty)
    )
    return wrap_tls(wrap_handshake(handshake_type, server_hello))


# 生成 Certificate 报文
def create_certificate(cert_path="server.crt"):
    try:
        with open(cert_path, 'rb') as f:
            cert_data = f.read()
        cert = x509.load_pem_x509_certificate(cert_data, default_backend())
        der_bytes = cert.public_bytes(Encoding.DER)
        global g_certificate
        g_certificate += len(der_bytes).__str__()
        write_string_to_file(g_certificate, "./1.txt", "w")
        cert_length = struct.pack('>I', len(der_bytes))[1:]
        cert_total_len = struct.pack('>I', len(der_bytes) + 3)[1:]
        certificates = (
                cert_total_len +
                cert_length +
                der_bytes
        )
        return wrap_tls(wrap_handshake(b'\x0B', certificates))
    except FileNotFoundError:
        print("Certificate file not found.")
        return None

def read_ecc_private_key_pem(file_path, password=None):
    """
    读取 ECC 私钥 PEM 文件
    :param file_path: 私钥 PEM 文件的路径
    :param password: 私钥的密码（如果有的话），默认为 None
    :return: ECC 私钥对象，如果读取失败则返回 None
    """
    try:
        with open(file_path, 'rb') as key_file:
            private_pem = key_file.read()
        private_key = serialization.load_pem_private_key(
            private_pem,
            password=password,
            backend=None
        )
        if isinstance(private_key, ec.EllipticCurvePrivateKey):
            print("成功读取 ECC 私钥")
            return private_key
        else:
            print("读取的不是 ECC 私钥")
    except FileNotFoundError:
        print(f"未找到私钥 PEM 文件: {file_path}")
    except ValueError:
        print("私钥密码错误或 PEM 文件格式错误")
    return None


def read_ecc_public_key_pem(file_path):
    """
    读取 ECC 公钥 PEM 文件
    :param file_path: 公钥 PEM 文件的路径
    :return: ECC 公钥对象，如果读取失败则返回 None
    """
    try:
        with open(file_path, 'rb') as key_file:
            public_pem = key_file.read()
        public_key = serialization.load_pem_public_key(
            public_pem,
            backend=None
        )
        if isinstance(public_key, ec.EllipticCurvePublicKey):
            print("成功读取 ECC 公钥")
            return public_key
        else:
            print("读取的不是 ECC 公钥")
    except FileNotFoundError:
        print(f"未找到公钥 PEM 文件: {file_path}")
    except ValueError:
        print("PEM 文件格式错误")
    return None

# 生成 ServerKeyExchange 报文
def create_server_key_exchange():
    private_key = read_ecc_private_key_pem("./ecc-key.pem")
    public_key = private_key.public_key()
    client_private_key = read_ecc_private_key_pem("./client-ecc-key.pem")
    client_public_key = client_private_key.public_key()
    shared_secret = private_key.exchange(ec.ECDH(), client_public_key)

    master_secret = HKDF(
        algorithm=hashes.SHA256(),
        length=48,
        salt=None,
        info=b'tls12 master secret',
    ).derive(shared_secret)
    # der_public = public_key.public_bytes(encoding=Encoding.DER, format=PublicFormat.SubjectPublicKeyInfo)

    der_public = public_key.public_bytes(Encoding.X962, PublicFormat.CompressedPoint)
    curve = ec.SECP256R1()
    uncompressed_public_key = ec.EllipticCurvePublicKey.from_encoded_point(curve, der_public)
    uncompressed_public_key_bytes = uncompressed_public_key.public_bytes(
        encoding=Encoding.X962,
        format=PublicFormat.CompressedPoint
    )
    isEqual = uncompressed_public_key_bytes == der_public
    curve_type = struct.pack('>H', 0x0017)
    der_public_len = struct.pack('>H', len(der_public))[1:]
    handshake = (
            b'\x03' + #椭圆曲线类型
            curve_type + # 命名曲线
            der_public_len + # 在 SECP256R1 曲线中，压缩公钥的长度为 33 字节，非压缩公钥的长度为 65 字节
            der_public # 服务器的椭圆曲线公钥
    )
    return wrap_tls(wrap_handshake(b'\x0C', handshake)), master_secret

class Test(TestCase):
    def test_create_server_key_exchange(self):
        create_server_key_exchange()

# 生成 ServerHelloDone 报文
def create_server_hello_done():
    return wrap_tls(wrap_handshake(b'\x0E', b''))


# 生成 ChangeCipherSpec 报文
def create_change_cipher_spec():
    record = b'\x14\x03\x03\x00\x01\x01'
    return record


# 生成 Finished 报文
def create_finished(handshake_messages, master_secret, is_client=False):
    label = b'server finished' if not is_client else b'client finished'
    verify_data = HKDF(
        algorithm=hashes.SHA256(),
        length=12,
        salt=b'',
        info=label,
    ).derive(master_secret)

    handshake = (
            b'\x14' +  # Handshake type: Finished
            struct.pack('>I', 12)[1:] +
            verify_data
    )
    return wrap_tls(handshake)


# 封装握手消息到 TLS 记录层
def wrap_tls(handshake_data):
    record_type = b'\x16'
    protocol_version = b'\x03\x03'
    record_length = struct.pack('>H', len(handshake_data))
    return record_type + protocol_version + record_length + handshake_data


def wrap_handshake(handshake_type: bytes, handshake_data: bytes):
    record_length = struct.pack('>I', len(handshake_data))[1:]
    return handshake_type + record_length + handshake_data





# 加密数据
def encrypt_data(data, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    encryptor = cipher.encryptor()
    padded_data = data + b'\x00' * ((16 - len(data) % 16) % 16)
    ciphertext = encryptor.update(padded_data) + encryptor.finalize()
    return ciphertext


# 解密数据
def decrypt_data(ciphertext, key, iv):
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
    decryptor = cipher.decryptor()
    padded_data = decryptor.update(ciphertext) + decryptor.finalize()
    return padded_data.rstrip(b'\x00')


# 发送加密数据
def send_encrypted_data(sock, data, key, iv):
    encrypted = encrypt_data(data, key, iv)
    record = (
            b'\x17' +  # Application data
            b'\x03\x03' +  # TLS 1.2
            struct.pack('>H', len(encrypted)) +
            encrypted
    )
    sock.send(record)


# 接收加密数据
def receive_encrypted_data(sock, key, iv):
    record_type = sock.recv(1)
    if record_type != b'\x17':
        raise ValueError("Not application data")

    version = sock.recv(2)
    length = struct.unpack('>H', sock.recv(2))[0]
    ciphertext = sock.recv(length)
    return decrypt_data(ciphertext, key, iv)


# 服务器端会话缓存
server_session_cache = {}


# 处理客户端 Hello 报文
def handle_client_hello(client_hello):
    session_id_length = client_hello[29]
    session_id = client_hello[30:30 + session_id_length]

    if session_id in server_session_cache:
        print("Session resumed:", session_id.hex())
        return server_session_cache[session_id]
    else:
        print("New session")
        session_id = uuid.uuid4().bytes
        master_secret = bytes(random.getrandbits(8) for _ in range(48))
        server_session_cache[session_id] = master_secret
        return master_secret


def server():
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )
    public_key = private_key.public_key()

    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, u"US"),
        x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"California"),
        x509.NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"My Company"),
        x509.NameAttribute(NameOID.COMMON_NAME, u"example.com"),
    ])
    cert = x509.CertificateBuilder().subject_name(
        subject
    ).issuer_name(
        issuer
    ).public_key(
        public_key
    ).serial_number(
        x509.random_serial_number()
    ).not_valid_before(
        datetime.datetime.utcnow()
    ).not_valid_after(
        datetime.datetime.utcnow() + datetime.timedelta(days=365)
    ).sign(private_key, hashes.SHA256(), default_backend())
    crt_name = "server.crt"
    if not os.path.exists(crt_name):
        with open(crt_name, "wb") as f:
            f.write(cert.public_bytes(Encoding.PEM))

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("0.0.0.0", 8888))
        sock.listen(1)
        print("Server listening on port 8888")
        conn, addr = sock.accept()
        print(f"Connection from {addr}")

        client_hello = conn.recv(4096)
        handle_client_hello(client_hello)

        server_hello = create_server_hello(client_hello)
        conn.send(server_hello)
        certificate = create_certificate()
        conn.send(certificate)
        server_key_exchange, master_secret= create_server_key_exchange()
        conn.send(server_key_exchange)
        server_hello_done = create_server_hello_done()

        if certificate is None:
            print("Failed to send certificate. Closing connection.")
            conn.close()
            return
        #
        server_responses = server_hello_done
        conn.send(server_responses)

        client_key_exchange = conn.recv(4096)
        client_change_cipher = conn.recv(4096)
        client_finished = conn.recv(4096)

        change_cipher_spec = create_change_cipher_spec()
        conn.send(change_cipher_spec)

        handshake_messages = server_responses + client_key_exchange
        server_finished = create_finished(handshake_messages, master_secret, is_client=False)
        conn.send(server_finished)

        client_random = client_hello[6:38]
        server_random = server_hello[6:38]
        client_write_key, server_write_key, client_write_iv, server_write_iv = derive_keys(
            master_secret, client_random, server_random
        )

        while True:
            try:
                data = receive_encrypted_data(conn, client_write_key, client_write_iv)
                if not data:
                    break
                print(f"Received from client: {data.decode()}")
                response = f"Server received: {data.decode()}"
                send_encrypted_data(conn, response.encode(), server_write_key, server_write_iv)
            except ValueError as e:
                print(f"Error: {e}")
                break


if __name__ == "__main__":
    server()
